/*
 * AES-128-CBC demo using openHiTLS EAL cipher API
 * Modes:
 *   --nist : NIST SP 800-38A F.2.1 single-block vector (check cipher equals expected)
 *   (default): 32-byte plaintext, PKCS#7 padding, check decrypt == plaintext
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stdbool.h>
#include "crypt_eal_cipher.h"
#include "bsl_sal.h"
#include "bsl_err.h"
#include "crypt_algid.h"
#include "crypt_errno.h"

void *StdMalloc(uint32_t len) { return malloc((size_t)len); }

void PrintLastError(void) {
    const char *file = NULL; uint32_t line = 0;
    BSL_ERR_GetLastErrorFileLine(&file, &line);
    printf("failed at file %s at line %u\n", file ? file : "(unknown)", line);
}

static void print_hex(const char *tag, const uint8_t *buf, uint32_t len) {
    printf("%s", tag);
    for (uint32_t i = 0; i < len; i++) printf("%02x", buf[i]);
    printf("\n");
}

int main(int argc, char **argv)
{
    bool use_nist = (argc > 1 && strcmp(argv[1], "--nist") == 0);

    /* NIST key/iv (valid for both modes) */
    uint8_t key[16] = {
        0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,
        0xab,0xf7,0x15,0x88,0x09,0xcf,0x4f,0x3c
    };
    uint8_t iv[16] = {
        0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,
        0x08,0x09,0x0a,0x0b,0x0c,0x0d,0x0e,0x0f
    };

    /* Plaintexts for two modes */
    uint8_t plain_nist[16] = {
        0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,
        0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a
    };
    uint8_t plain_default[32] = {
        0x6b,0xc1,0xbe,0xe2,0x2e,0x40,0x9f,0x96,
        0xe9,0x3d,0x7e,0x11,0x73,0x93,0x17,0x2a,
        0xae,0x2d,0x8a,0x57,0x1e,0x03,0xac,0x9c,
        0x9e,0xb7,0x6f,0xac,0x45,0xaf,0x8e,0x51
    };

    /* NIST expected cipher for single-block only */
    uint8_t expect_cipher_nist[16] = {
        0x76,0x49,0xab,0xac,0x81,0x19,0xb2,0x46,
        0xce,0xe9,0x8e,0x9b,0x12,0xe9,0x19,0x7d
    };

    const uint8_t *plain = use_nist ? plain_nist : plain_default;
    uint32_t plainLen = use_nist ? sizeof(plain_nist) : sizeof(plain_default);

    uint8_t cipher[96] = {0};
    uint8_t decrypted[96] = {0};
    uint32_t outLen = sizeof(cipher);
    uint32_t outTotal = 0;
    int32_t ret;

    print_hex("plain text to be encrypted: ", plain, plainLen);

    /* Create AES-128-CBC context */
    CRYPT_EAL_CipherCtx *ctx = CRYPT_EAL_CipherNewCtx(CRYPT_CIPHER_AES128_CBC);
    if (ctx == NULL) { PrintLastError(); BSL_ERR_DeInit(); return 1; }

    /* ---- Encrypt ---- */
    ret = CRYPT_EAL_CipherInit(ctx, key, sizeof(key), iv, sizeof(iv), true);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }

    ret = CRYPT_EAL_CipherSetPadding(ctx, CRYPT_PADDING_PKCS7);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }

    outLen = sizeof(cipher);
    ret = CRYPT_EAL_CipherUpdate(ctx, plain, plainLen, cipher, &outLen);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }
    outTotal += outLen;

    outLen = sizeof(cipher) - outTotal;
    ret = CRYPT_EAL_CipherFinal(ctx, cipher + outTotal, &outLen);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }
    outTotal += outLen;

    print_hex("cipher text value is: ", cipher, outTotal);

    if (use_nist) {
        /* In NIST single-block mode, there should be exactly 16 bytes output (no extra padding block) ONLY IF
           we disable padding. But since we use PKCS#7, an extra block would be added for block-multiple inputs.
           Therefore, in NIST mode we purposely feed 16B plain and still accept PKCS#7 extra block.
           To match NIST exactly, compare only the first block with expected. */
        if (outTotal < 16 || memcmp(cipher, expect_cipher_nist, 16) != 0) {
            printf("ciphertext mismatch with NIST single-block (first block)!\n");
            goto EXIT;
        }
    }

    /* ---- Decrypt ---- */
    /* Re-init for decryption (IV must be reset to the same bytes as used for encryption) */
    ret = CRYPT_EAL_CipherInit(ctx, key, sizeof(key), iv, sizeof(iv), false);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }

    ret = CRYPT_EAL_CipherSetPadding(ctx, CRYPT_PADDING_PKCS7);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }

    uint32_t decTotal = 0;
    outLen = sizeof(decrypted);
    ret = CRYPT_EAL_CipherUpdate(ctx, cipher, outTotal, decrypted, &outLen);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }
    decTotal += outLen;

    outLen = sizeof(decrypted) - decTotal;
    ret = CRYPT_EAL_CipherFinal(ctx, decrypted + decTotal, &outLen);
    if (ret != CRYPT_SUCCESS) { printf("error code is %x\n", ret); PrintLastError(); goto EXIT; }
    decTotal += outLen;

    print_hex("decrypted plain text value is: ", decrypted, decTotal);

    if (decTotal != plainLen || memcmp(decrypted, plain, plainLen) != 0) {
        printf("plaintext comparison failed\n");
        goto EXIT;
    }

    printf("AES-128-CBC demo PASS\n");
    ret = 0;

EXIT:
    CRYPT_EAL_CipherFreeCtx(ctx);
    BSL_ERR_DeInit();
    return ret;
}